from common.variables import time_counts
from datetime import datetime
from trade.broker.base import BrokerBase
from trade.strategy import BaseStrategy


class BackTestModel:
    broker: BrokerBase
    strategy: BaseStrategy

    def __init__(self, broker, strategy):
        strategy.set_broker_model(broker)
        self.broker = broker
        self.strategy = strategy

    def _run_sliding(self):
        try:
            while self.broker.slide():
                if (unimrr := self.broker.get_unimrr()) < 2:
                    print(f"爆仓了， {unimrr}")
                    self.broker.force_clear_positions()
                    self.broker.force_clear_balances()
                    self.broker.force_clear_positions()
                    self.broker.force_clear_balances()
                    raise Exception(f"爆仓了，赶紧强平")

                self.strategy.apply()
                print(
                    f"{self.broker.get_cur_time()} 当前比特币：{self.broker.get_data('BTC/USDT', 'close', 0)}， "
                    f"当前资产总额： {self.broker.get_account_value()}, 当前unimrr: {unimrr}"
                )

        finally:

            self.strategy.after_all()

    def run_backtest(self, start_date: str, end_date: str):
        self.broker.reset()
        pre_start = (
            datetime.fromisoformat(start_date).timestamp()
            - self.strategy.preloads * time_counts[self.strategy.timeframe]
        )
        pre_start = datetime.fromtimestamp(pre_start).isoformat()
        self.broker.init_get_symbol_and_factors(
            pre_start,
            start_date,
            self.strategy.need_symbols,
            self.strategy.need_bfnames,
            self.strategy.timeframe,
        )
        self.broker.prepare_sliding(
            start_date,
            end_date,
            self.strategy.need_symbols,
            self.strategy.need_bfnames,
            self.strategy.timeframe,
        )
        self._run_sliding()

    def run_real(self):
        start_date = datetime.now().isoformat()
        pre_start = (
            datetime.fromisoformat(start_date).timestamp()
            - self.strategy.preloads * time_counts[self.strategy.timeframe]
        )
        self.broker.init_get_symbol_and_factors(
            pre_start,
            start_date,
            self.strategy.need_symbols,
            self.strategy.need_bfnames,
            self.strategy.timeframe,
        )
        self.broker.prepare_sliding(
            start_date,
            "",
            self.strategy.need_symbols,
            self.strategy.need_bfnames,
            self.strategy.timeframe,
        )
        self._run_sliding()
